package ssau.knyazev.modules;

import java.io.*;

import ssau.knyazev.common.CudaCRuntimeException;
import ssau.knyazev.common.UniException;
import jcuda.CudaException;
import jcuda.driver.*;

public class CudaCompiler {
	
	public static String preparePtxFile(String cuFileName) throws CudaCRuntimeException {
		return preparePtxFile(cuFileName, false);
	}
	
	public static String preparePtxFile(String cuFileName, boolean recompile) throws CudaCRuntimeException {
		return preparePtxFile(cuFileName, recompile, "");
	}
	
	public static String preparePtxFile(String cuFileName, boolean recompile, String nvccParams) throws CudaCRuntimeException {
		try{
			int endIndex = cuFileName.lastIndexOf('.');
			if (endIndex == -1) {
				endIndex = cuFileName.length() - 1;
			}
			String ptxFileName = cuFileName.substring(0, endIndex + 1) + "ptx";
			File ptxFile = new File(ptxFileName);
			if (!recompile && ptxFile.exists()) {
				return ptxFileName; 
			}
			File cuFile = new File(cuFileName);
			if (!cuFile.exists()) {
				throw new IOException("Input file not found: " + cuFileName);
			}
			String modelString = "-m" + System.getProperty("sun.arch.data.model") + "" + nvccParams;
			String command = "nvcc " + modelString + " -ptx " + cuFile.getPath() + " -o " + ptxFileName;
			System.out.println("Executing\n" + command);
			Process process = Runtime.getRuntime().exec(command);
			String errorMessage = new String(toByteArray(process.getErrorStream()));
			String outputMessage = new String(toByteArray(process.getInputStream()));
			int exitValue = 0;
			try {
				exitValue = process.waitFor();
			} catch (InterruptedException e) {
				Thread.currentThread().interrupt();
				throw new UniException("Interrupted while waiting for nvcc output", e);
			}
			if (exitValue != 0) {
				System.out.println("nvcc process exitValue " + exitValue);
				System.out.println("errorMessage:\n" + errorMessage);
				System.out.println("outputMessage:\n" + outputMessage);
				throw new IOException("Could not create .ptx file: " + errorMessage);
			}
			System.out.println("Finished creating PTX file");
			return ptxFileName;
		} catch (Exception e){
			throw new CudaCRuntimeException(e);
		}
	}

	private static byte[] toByteArray(InputStream inputStream) throws IOException {
		ByteArrayOutputStream baos = new ByteArrayOutputStream();
		byte buffer[] = new byte[8192];
		while (true) {
			int read = inputStream.read(buffer);
			if (read == -1) {
				break;
			}
			baos.write(buffer, 0, read);
		}
		return baos.toByteArray();
	}
	
	public static CUfunction getFunction(String functionName, String moduleName) throws CudaCRuntimeException{
		try{
			CUmodule module = new CUmodule();
			JCudaDriver.cuModuleLoad(module, moduleName);
			CUfunction function = new CUfunction();
			JCudaDriver.cuModuleGetFunction(function, module, functionName);
			return function;
		} catch (CudaException e){
			throw new CudaCRuntimeException("Function with name <" + functionName + "> does not exist in module <" + moduleName + ">");
		}
		
	}
	
}
